/*
 * Copyright 2022 George Bartolomey Licensed under
 * the Apache License, Version 2.0 (the «License»);
*/

#include <jack/jack.h>
#include <jack/midiport.h>
#include <stdio.h>
#include <signal.h>
#include <stdlib.h>
#include <unistd.h>
#include <getopt.h>
#include <string.h>
#include <fftw3.h>
#include <math.h>

#define DEFAULT_BUFFER_SIZE 4096
#define DEFAULT_AMP_THRESHOLD 20
#define DEFAULT_ACC_THRESHOLD 0.02
#define DEFAULT_EVENT_RANGE_START 0
#define DEFAULT_EVENT_RANGE_END 255
#define DEFAULT_CLIENT_NAME "w2midi"

#define CLIENT_NAME_LENGTH 32
#define M_PI 3.14159265359
#define HANN_a0 0.53836
#define HANN_a1 0.46164
#define NOTES_BITMAP_SIZE 256

/*Спец. состояние библиотеки FFTW*/
fftwf_plan main_fftwf_plan;
fftwf_complex *fft_out;
jack_client_t *client;
jack_port_t *input_port;
jack_port_t *output_port;
jack_nframes_t sample_rate;
/*Значение оконной функции*/
float *window;
/*Буфер, содержащий предыдущий сигнал и текущий; используется ПФ*/
float *buffer;
int solo_mode;
int pitch_mode;
unsigned int buffer_size;
/*Таблица полутон-диск. частота из ПФ*/
unsigned short prev_pitch;
unsigned short *pitch_table;
unsigned char *semitone_bins;
/*Громкости полутонов*/ unsigned char *note_pressures;
unsigned char *prev_note_pressures;

/*Настройки*/
float amp_threshold;
float acc_threshold;
unsigned int event_range_start;
unsigned int event_range_end;
char client_name[CLIENT_NAME_LENGTH];

char *help_text =
#include "help.txt"
;

void init();
void parse_args(int argc, char **argv);
/*Рабочий цикл*/
int process(jack_nframes_t nframes, void *arg);
/*Вычислить окно Хемминга*/
void generate_window(float *out, int n);
/*Вычислить таблицу частота-полутон*/
void generate_semitone_bins();
/*Преобразовать спектр частот в спектр нот*/
void calculate_note_pressures(float *spectral_density);
/*Убрать обертоны*/
void remove_obertones();
void remove_all_other();
void send_events(void *output_buffer);
void send_pitch(void *output_buffer, float *spectral_density);
void array_mul(float *a, float *b, int n);
/*|z|^2 Преобразовать комплексные амлитуды в спектральную плотность*/
void cmplx_amp_to_spectral_density(fftwf_complex *in, int n);

int main(int argc, char **argv) {
    solo_mode = 0;
    pitch_mode = 0;
    buffer_size = DEFAULT_BUFFER_SIZE;
    amp_threshold = DEFAULT_AMP_THRESHOLD;
    acc_threshold = DEFAULT_ACC_THRESHOLD;
    event_range_start = DEFAULT_EVENT_RANGE_START;
    event_range_end = DEFAULT_EVENT_RANGE_END;
    strcpy(client_name, DEFAULT_CLIENT_NAME);
    parse_args(argc, argv);
    init();

    if (jack_activate(client)) {
        fprintf(stderr, "Failed to activate JACK client!\n");
        exit(EXIT_FAILURE);
    }
    sleep(-1);
    jack_client_close(client);
    return 0;
}

void parse_args(int argc, char **argv) {
    int opt;
    while ((opt = getopt(argc, argv, "b:d:a:n:s:e:n:hop")) != -1) {
        switch (opt) {
        case 'b':
            buffer_size = atoi(optarg);
            break;
        case 'd':
            amp_threshold = atof(optarg);
            break;
        case 'a':
            acc_threshold = atof(optarg);
            break;
        case 'n':
            strncpy(client_name, optarg, CLIENT_NAME_LENGTH);
            break;
        case 's':
            event_range_start = atoi(optarg);
            break;
        case 'e':
            event_range_end = atoi(optarg);
            break;
        case 'o':
            solo_mode = 1;
            break;
        case 'p':
            printf("pitch\n");
            pitch_mode = 1;
            break;
        case 'h':
            printf("%s", help_text);
            exit(0);
        }
    }
    if ((buffer_size & (buffer_size - 1)) != 0 && buffer_size > 8) {
        fprintf(stderr, "Buffer size must be power of two "
                "and greater than 8\n");
        exit(1);
    }
    if (event_range_start < 0 || event_range_end > 255
            || event_range_start > event_range_end) {
        fprintf(stderr, "Notes range must be from 0 to 255");
        exit(1);
    }
}

void init() {
    client = jack_client_open(client_name, JackNullOption, NULL);
    if (!client) {
        fprintf(stderr, "Failed to connect to JACK server!\n");
        exit(EXIT_FAILURE);
    }
    jack_set_process_callback(client, process, 0);
    input_port = jack_port_register(client, "in", JACK_DEFAULT_AUDIO_TYPE,
            JackPortIsInput, buffer_size / 2);
    output_port = jack_port_register(client, "out", JACK_DEFAULT_MIDI_TYPE,
            JackPortIsOutput, 0);
    if (!input_port || !output_port) {
        fprintf(stderr, "Failed to register JACK ports!\n");
        exit(EXIT_FAILURE);
    }
    jack_set_buffer_size(client, buffer_size / 2);
    sample_rate = jack_get_sample_rate(client);
    
    buffer = (float*)fftwf_malloc(sizeof(float) * buffer_size);
    window = (float*)malloc(sizeof(float) * buffer_size);
    semitone_bins = (unsigned char*)malloc(buffer_size / 2);
    note_pressures = (unsigned char*)malloc(NOTES_BITMAP_SIZE);
    prev_note_pressures = (unsigned char*)malloc(NOTES_BITMAP_SIZE);
    pitch_table = (unsigned short*)malloc(buffer_size / 2 * sizeof(unsigned short));
    if (!buffer || !window || !semitone_bins
            || !note_pressures || !prev_note_pressures || !pitch_table) {
        fprintf(stderr, "Failed to allocate memory!\n");
        exit(EXIT_FAILURE);
    }
    fft_out = (fftwf_complex*)fftwf_malloc(sizeof(fftwf_complex) * buffer_size);
    main_fftwf_plan = fftwf_plan_dft_r2c_1d(buffer_size, buffer,
            fft_out, FFTW_ESTIMATE);

    generate_window(window, buffer_size);
    memset(prev_note_pressures, 0, NOTES_BITMAP_SIZE);
    memset(buffer, 0, buffer_size * sizeof(float));
    prev_pitch = 0;
    generate_semitone_bins();
}


int process(jack_nframes_t nframes, void *arg) {
    jack_default_audio_sample_t *input_buffer;
    void* output_buffer;
    float *spectral_density;

    if (nframes <= 0) {
        return 0;
    }
    input_buffer = jack_port_get_buffer(input_port, nframes);
    output_buffer = jack_port_get_buffer(output_port, nframes);
    jack_midi_clear_buffer(output_buffer);
    memset(note_pressures, 0, NOTES_BITMAP_SIZE);

    memcpy(&buffer[buffer_size / 2], input_buffer, 
            buffer_size / 2 * sizeof(float));
    array_mul(buffer, window, buffer_size);
    fftwf_execute(main_fftwf_plan);
    memcpy(buffer, input_buffer, buffer_size / 2 * sizeof(float));
    cmplx_amp_to_spectral_density(fft_out, buffer_size / 2);
    spectral_density = (float*)fft_out;
    calculate_note_pressures(spectral_density);
    if (pitch_mode)
        send_pitch(output_buffer, spectral_density);
    if (solo_mode)
        remove_all_other();
    else
        remove_obertones();
    send_events(output_buffer);
    memcpy(prev_note_pressures, note_pressures, NOTES_BITMAP_SIZE);
    return 0;    
}

/* Hann function window */
void generate_window(float *out, int n) {
    int i;
    float c;
    c = 2 * M_PI / (float)(n - 1);
    for (i = 0; i < n; i++) {
        out[i] = HANN_a0 - HANN_a1 * cos(c * (float)i); 
    }
}

void array_mul(float *a, float *b, int n) {
    int i;
    for (i = 0; i < n; i++) {
        a[i] *= b[i];
    }
}

/*|z|^2*/
void cmplx_amp_to_spectral_density(fftwf_complex *in, int n) {
    int i;
    float *out;
    out = (float*)in;
    for (i = 0; i < n; i++) {
        out[i] = in[i][0]*in[i][0] + in[i][1]*in[i][1];
    }
}

void generate_semitone_bins() {
    int i;
    float freq;
    float semitone;
    float difference;
    semitone_bins[0] = 0;
    for (i = 0; i < buffer_size / 2; i++) {
        freq = (float)i * sample_rate / (float)buffer_size;
        semitone = 12 * log2(freq / 440) + 69;
        difference = abs(round(semitone) - semitone);
        if (semitone > 0 && semitone < 256 && difference < acc_threshold) {
            semitone_bins[i] = round(semitone);
            if (pitch_mode)
                pitch_table[i] = ((round(semitone) - semitone) * 0x2000 / 2) + 0x2000;
        }
        else {
            semitone_bins[i] = 0;
            if (pitch_mode)
                pitch_table[i] = 0;
        }
    }
}

void calculate_note_pressures(float *spectral_density) {
    int i;
    float db;
    for (i = event_range_start; i < buffer_size / 2 && i < event_range_end; i++) {
        db = 10 * log10(spectral_density[i]);
        if (db > amp_threshold) {
            note_pressures[semitone_bins[i]] = (db < 256 ? db : 255);
        }
    }
}

void remove_obertones() {
    int i;
    for (i = 0; i < NOTES_BITMAP_SIZE - 3; i++) {
        if (note_pressures[i] < note_pressures[i - 1] ||
            note_pressures[i] < note_pressures[i - 2] ||
                note_pressures[i] < note_pressures[i + 1] ||
            note_pressures[i] < note_pressures[i + 2]) {
            note_pressures[i] = 0;
        }        
    }
}

void remove_all_other() {
    int i;
    int max_i;
    unsigned char max;
    max_i = 0;
    max = 0;
    for (i = 0; i < NOTES_BITMAP_SIZE - 3; i++) {
        if (note_pressures[i] > max) {
            max_i = i;
            max = note_pressures[i];
        }
    }
    memset(note_pressures, 0, NOTES_BITMAP_SIZE);
    note_pressures[max_i] = max;
}

void send_pitch(void *output_buffer, float *spectral_density) {
    int i;
    int max_i;
    float max;
    unsigned short pitch;
    unsigned char *midi;
    max_i = -1;
    max = -INFINITY;
    for (i = 1; i < buffer_size / 2; i++) {
        if (max < spectral_density[i]) {
            max = spectral_density[i];
            max_i = i;
        }
    }
    pitch = pitch_table[max_i];
        midi = jack_midi_event_reserve(output_buffer, 0, 3);
        midi[0] = 0xE0;
        midi[1] = pitch & 0xFF;
        midi[2] = pitch >> 8 & 0xFF;
        printf("%i\n", pitch);
    prev_pitch = pitch;
}

void send_events(void *output_buffer) {
    int i;
    unsigned char *midi;
    unsigned int pressure;
    i = event_range_start;
    for (i = event_range_start; i < event_range_end; i++) {
        if (prev_note_pressures[i] == 0 && note_pressures[i] > 0) {
            midi = jack_midi_event_reserve(output_buffer, i, 3);
            midi[0] = 0x90;
            midi[1] = i;
            pressure = note_pressures[i] * 32 / amp_threshold;
            midi[2] = pressure < 100 ? pressure : 100;
        }
        else if (prev_note_pressures[i] > 0 && note_pressures[i] == 0) {
            midi = jack_midi_event_reserve(output_buffer, i, 3);
            midi[0] = 0x80;
            midi[1] = i;
            midi[2] = 64;
        }
    }
}

